#!/usr/bin/env python3
"""
Individual PBPK Modeling Script

This script uses existing modules in the src folder to implement PBPK modeling with
different absorption and distribution modes:
- Absorption modes: first_order, adam
- Distribution modes: minimal_pbpk, full_pbpk

Author: BlackCat@Cppo
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# Add src path to system path
# sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../src'))
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
print(f"Project root path: {project_root}")
sys.path.insert(0, project_root)

from src.absorption.absorption_module import AbsorptionModel
from src.absorption.adam_model import ADAMAbsorptionModel
from src.distribution.distribution_module import DistributionModel
from src.elimination.elimination_module import EliminationModel
from src.elimination.inhibition_clearance import Inhibition_Model
from src.elimination.inhibition_clearance import InhibitionModelIntegrator


class IndividualPBPKSimulator:
    """Individual PBPK Simulator"""

    def __init__(self):
        """Initialize simulator"""
        self.elimination_model = None
        self.default_params = self._get_default_parameters()
        self.time_points = np.linspace(0, 24, 48)  # 24 hours (1 day), 48 time points

    def _get_default_parameters(self):
        """Get default parameters"""
        return {
            # Basic parameters
            'dose': 100.0,  # mg
            'body_weight': 70.0,  # kg

            # Absorption parameters
            'ka': 1.0,  # 吸收速率常数 (1/h)
            'fa': 0.9,  # 吸收分数
            'fg': 0.95,  # 肠道代谢分数
            'fh': 0.90,  # 肝脏代谢分数

            # Distribution parameters
            'Vp': 3.0,  # 血浆容积 (L)
            'Ve': 12.0,  # 细胞外液容积 (L)
            'Vt': 12.0,  # 组织容积 (L)
            'fu': 0.1,  # 游离分数
            'BP': 1.0,  # 血液-血浆浓度比

            # 组织分配系数
            'Kp_liver': 2.0,
            'Kp_kidney': 2.5,
            'Kp_brain': 0.5,
            'Kp_lung': 1.5,
            'Kp_heart': 1.8,
            'Kp_muscle': 1.2,
            'Kp_adipose': 0.8,

            # Elimination parameters (realistic for 70kg adult)
            'CL_hepatic': 2.0,  # 肝脏清除率 (L/h)
            'CL_renal': 1.5,  # 肾脏清除率 (L/h)
            'CL_biliary': 0.5,  # 胆汁清除率 (L/h)
            'CL_other': 0.2,  # 其他清除率 (L/h)

            # Enzyme kinetic parameters (核心动力学参数)
            'Vmax': 100.0,  # 酶促反应最大速率 (μmol/min)
            'Km': 10.0,  # 米氏常数 (μg/mL)
            'Ki': 1.0,  # 抑制常数，默认1.0 (μg/mL)
            'Clint': 5.0,  # 内在清除率 (L/h)
            'fu_mic': 1.0,  # 微粒体游离分数，默认1.0

            # ADAM model parameters
            'use_adam': False,
            'solubility': 1.0,  # mg/mL
            'permeability': 1e-4,  # cm/s
            'particle_size': 50.0,  # μm
            'density': 1.2,  # g/cm³
            'molecular_weight': 300.0,  # g/mol
            'pKa': 7.4,
            'logP': 2.0
        }

    def run_absorption(self, model_type, params, time_points):
        """Run absorption model"""
        if model_type == 'first_order':
            model = AbsorptionModel(params)
            result = model.calculate_absorption(params, time_points)
            return {
                'plasma_concentration': result['plasma_concentration'],
                'cumulative_absorption': result['cumulative_absorption'],
                'absorption_rate': result['absorption_rate'],
                'bioavailability': result['bioavailability']
            }

        elif model_type == 'adam':
            adam_model = ADAMAbsorptionModel()
            result = adam_model.calculate_adam_absorption(time_points, params)
            return {
                'plasma_concentration': result.get('plasma_concentration', result['cumulative_absorption']),
                'cumulative_absorption': result['cumulative_absorption'],
                'absorption_rate': result['total_absorption_rate'],
                'bioavailability': result['bioavailability']
            }

        else:
            raise ValueError(f"Unknown absorption model type: {model_type}")

    def run_distribution(self, model_type, params, time_points, absorption_result):
        """Run distribution model"""
        absorption_input = {
            'dose': params['dose'],
            'cumulative': absorption_result['cumulative_absorption'],
            'plasma_concentration': absorption_result['plasma_concentration'],
            'bioavailability': absorption_result['bioavailability']
        }

        model = DistributionModel(params)

        if model_type == 'minimal_pbpk':
            # Simplified PBPK distribution model (一室模型)
            result = model.calculate_distribution(params, time_points, absorption_input)
            return {
                'plasma_concentration': result['plasma_concentration'],
                'distribution_volume': result['distribution_volume'],
                'tissue_concentrations': result['tissue_concentrations']
            }

        elif model_type == 'full_pbpk':
            # Full PBPK distribution model (多室模型)
            # 使用新添加的多室模型函数
            multi_comp_result = model.calculate_multi_compartment_plasma_concentration(
                params, time_points, absorption_input
            )

            # 计算分布容积
            tissue_partition = model.calculate_tissue_partition_coefficients(params)
            distribution_volume = model.calculate_V_d(params, tissue_partition)

            # 计算脑部分布
            brain_result = model.calculate_brain_distribution(
                params, time_points, multi_comp_result['plasma_concentration']
            )

            # 合并组织浓度
            tissue_concentrations = multi_comp_result.get('tissue_concentrations', {})
            tissue_concentrations['brain'] = brain_result['concentration']

            return {
                'plasma_concentration': multi_comp_result['plasma_concentration'],
                'distribution_volume': distribution_volume,
                'tissue_concentrations': tissue_concentrations,
                'brain_concentration': brain_result['concentration'],
                'brain_partition': brain_result['partition_coefficient']
            }

        else:
            raise ValueError(f"Unknown distribution model type: {model_type}")

    def run_elimination(self, model_type, params, time_points, distribution_result):
        """Run elimination model"""
        model = EliminationModel(params)

        # 对于完整PBPK模型，分布结果已经包含消除过程
        # 对于最小模型，需要计算消除
        if model_type == 'full_pbpk':
            # 直接使用完整PBPK模型的血浆浓度
            plasma_conc = distribution_result['plasma_concentration']
        else:
            # 计算消除过程
            elimination_result = model.calculate_elimination(params, time_points, distribution_result)
            plasma_conc = elimination_result.get('plasma_concentration', distribution_result['plasma_concentration'])

        # 计算清除率
        try:
            hepatic_clearance = model.calculate_hepatic_clearance(params)
            renal_clearance = model.calculate_renal_clearance(params)
            biliary_clearance = model.calculate_biliary_clearance(params)

            # 处理可能的字典返回情况
            if isinstance(hepatic_clearance, dict):
                hepatic_clearance = hepatic_clearance.get('hepatic_clearance', 2.0)
            if isinstance(renal_clearance, dict):
                renal_clearance = renal_clearance.get('total_renal_clearance', 1.5)
            if isinstance(biliary_clearance, dict):
                biliary_clearance = biliary_clearance.get('biliary_clearance', 0.5)

            total_clearance = hepatic_clearance + renal_clearance + biliary_clearance
        except:
            # 使用默认值
            hepatic_clearance = 2.0
            renal_clearance = 1.5
            biliary_clearance = 0.5
            total_clearance = 4.0

        return {
            'plasma_concentration': plasma_conc,
            'clearance_rates': {
                'hepatic': hepatic_clearance,
                'renal': renal_clearance,
                'biliary': biliary_clearance,
                'total': total_clearance
            },
            'eliminated_amount': np.zeros_like(time_points)  # 简化处理
        }

    def run_complete_pbpk(self, absorption_model, distribution_model, inhibition_model, substrate_params=None, inhibitor_params=None):
        """Run complete PBPK simulation"""
        if substrate_params is None:
            substrate_params = self.default_params.copy()
        if inhibitor_params is None:
            inhibitor_params = self.default_params.copy()

        print(f"\nRunning simulation: absorption={absorption_model}, distribution={distribution_model}")

        # 1. 吸收阶段
        substrate_absorption_result = self.run_absorption(absorption_model, substrate_params, self.time_points)
        print(f"  Absorption completed: bioavailability={substrate_absorption_result['bioavailability']:.2%}")
        inhibitor_absorption_result = self.run_absorption(absorption_model, inhibitor_params, self.time_points)
        print(f"  Absorption completed: bioavailability={inhibitor_absorption_result['bioavailability']:.2%}")

        # 2. 分布阶段
        substrate_distribution_result = self.run_distribution(distribution_model, substrate_params, self.time_points, substrate_absorption_result)
        print(f"  Distribution completed: volume of distribution={substrate_distribution_result['distribution_volume']:.2f} L")
        inhibitor_distribution_result = self.run_distribution(distribution_model, inhibitor_params, self.time_points, inhibitor_absorption_result)
        print(f"  Distribution completed: volume of distribution={inhibitor_distribution_result['distribution_volume']:.2f} L")

        # 3. 消除阶段 - 修复抑制计算
        # 创建消除模型实例
        elimination_model = EliminationModel(substrate_params)
        
        # 计算基础消除（无抑制剂）
        base_elimination = elimination_model.calculate_elimination(
            substrate_params, self.time_points, substrate_distribution_result
        )
        
        # 准备抑制计算参数（合并底物和抑制剂参数）
        inhibition_params = substrate_params.copy()
        inhibition_params.update({
            'inhibitor_concentration': np.mean(inhibitor_distribution_result['plasma_concentration']),
            'substrate_concentration': np.mean(substrate_distribution_result['plasma_concentration'])
        })
        
        # 创建抑制剂积分器
        eliminator = InhibitionModelIntegrator(elimination_model)

        # 应用抑制效应到清除率计算
        inhibition_result = eliminator.apply_inhibition_to_clearance(
            inhibition_mode=inhibition_model,
            parameters=inhibition_params,
            substrate_concentration=substrate_distribution_result['plasma_concentration'],
            inhibitor_concentration=inhibitor_distribution_result['plasma_concentration']
        )

        # 计算抑制后的血浆浓度时间曲线
        inhibited_clearance_rates = base_elimination['clearance_rates'].copy()
        
        # 更新肝脏清除率（主要受CYP抑制影响）
        if 'hepatic_clearance' in inhibition_result:
            old_hepatic = inhibited_clearance_rates['hepatic']
            new_hepatic = inhibition_result['hepatic_clearance']
            
            # 处理数组或标量情况
            if isinstance(new_hepatic, np.ndarray):
                new_hepatic_scalar = float(np.mean(new_hepatic))
            else:
                new_hepatic_scalar = float(new_hepatic)
                
            inhibited_clearance_rates['hepatic'] = max(new_hepatic_scalar, 0.01)  # 防止清除率为0
            inhibited_clearance_rates['total'] = (
                inhibited_clearance_rates['hepatic'] + 
                inhibited_clearance_rates['renal'] + 
                inhibited_clearance_rates['biliary'] + 
                inhibited_clearance_rates['other']
            )
            
            print(f"    Hepatic clearance: {old_hepatic:.2f} → {inhibited_clearance_rates['hepatic']:.2f} L/h")
            
            # Handle inhibition ratio formatting
            inhibition_ratio = inhibition_result.get('inhibition_ratio', 0.0)
            if isinstance(inhibition_ratio, np.ndarray):
                inhibition_ratio_scalar = float(np.mean(inhibition_ratio))
            else:
                inhibition_ratio_scalar = float(inhibition_ratio)
        else:
            inhibition_ratio_scalar = 0.0
            
        # 重新计算消除半衰期和血浆浓度
        inhibited_half_life = elimination_model.calculate_elimination_half_life(
            substrate_params, inhibited_clearance_rates, substrate_distribution_result
        )

        # 使用抑制后的清除率重新计算血浆浓度
        inhibited_plasma_concentration = self._calculate_inhibited_concentration(
            substrate_params, self.time_points, substrate_distribution_result, inhibited_clearance_rates
        )

        # 合并消除结果
        elimination_result = {
            'plasma_concentration': inhibited_plasma_concentration,
            'clearance_rates': inhibited_clearance_rates,
            'inhibition_ratio': inhibition_ratio_scalar,
            'elimination_half_life': inhibited_half_life,
            'base_elimination': base_elimination,
            'inhibition_details': inhibition_result
        }
        print(f"  Elimination completed: total clearance={inhibited_clearance_rates['total']:.2f} L/h")

        # 返回完整结果
        return {
            'time_points': self.time_points,
            'substrate_absorption': substrate_absorption_result,
            'inhibition_absorption': inhibitor_absorption_result,
            'substrate_distribution': substrate_distribution_result,
            'inhibition_distribution': inhibitor_distribution_result,
            'elimination': elimination_result,
            'models': {
                'absorption': absorption_model,
                'distribution': distribution_model,
                'inhibition': inhibition_model,
            }
        }

    def _calculate_inhibited_concentration(self, params, time_points, distribution_result, clearance_rates):
        """计算抑制后的血浆浓度时间曲线"""
        # 获取分布容积
        if 'distribution_volume' in distribution_result:
            V_d = distribution_result['distribution_volume']
        else:
            V_d = params.get('V_d', 15.0)
        
        # 获取总清除率
        total_clearance = clearance_rates['total']
        
        # 获取吸收阶段的参数
        dose = params.get('dose', 100.0)
        ka = params.get('ka', 1.0)
        
        # 计算消除速率常数
        ke = total_clearance / V_d
        
        # 计算生物利用度
        F = params.get('fa', 0.9) * params.get('fg', 0.95) * params.get('fh', 0.85)
        
        # 使用一室模型计算浓度-时间曲线
        # C(t) = (F * Dose * ka) / (V_d * (ka - ke)) * (e^(-ke*t) - e^(-ka*t))
        
        # 防止数值问题
        if abs(ka - ke) < 1e-6:
            ke = ke * 1.01
            
        # 计算浓度
        concentration = (F * dose * ka) / (V_d * abs(ka - ke)) * (
            np.exp(-ke * time_points) - np.exp(-ka * time_points)
        )
        
        # 确保非负
        concentration = np.maximum(concentration, 0.0)
        
        return concentration

    def calculate_pk_metrics(self, time_points, concentration):
        """Calculate pharmacokinetic parameters"""
        # C_max and T_max
        C_max = np.max(concentration)
        T_max = time_points[np.argmax(concentration)]

        # AUC (trapezoidal method)
        AUC = np.trapezoid(concentration, time_points)

        # Terminal half-life
        if len(concentration) >= 10:
            terminal_points = concentration[-10:]
            if np.any(terminal_points > 0):
                log_conc = np.log(terminal_points[terminal_points > 0])
                time_term = time_points[-len(terminal_points):][terminal_points > 0]
                if len(log_conc) > 1:
                    slope, _ = np.polyfit(time_term, log_conc, 1)
                    t_half = -0.693 / slope if slope < 0 else 999
                else:
                    t_half = 999
            else:
                t_half = 999
        else:
            t_half = 999

        return {
            'C_max': C_max,
            'T_max': T_max,
            'AUC': AUC,
            't_half': t_half
        }

    def plot_all_models(self, results_dict):
        """Plot concentration-time curves for all models in separate plots"""
        # Create a 2x2 grid of subplots
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        axes = axes.flatten()

        # Define model combinations and their order
        model_combinations = [
            ('first_order', 'minimal_pbpk', Inhibition_Model.COMPETITIVE_INHIBITION),
            ('first_order', 'full_pbpk', Inhibition_Model.COMPETITIVE_INHIBITION),
            ('adam', 'minimal_pbpk', Inhibition_Model.COMPETITIVE_INHIBITION),
            ('adam', 'full_pbpk', Inhibition_Model.COMPETITIVE_INHIBITION)
        ]

        colors = ['blue', 'red', 'green', 'purple']

        # Plot each model in its own subplot
        for idx, (abs_model, dist_model, inhib_model) in enumerate(model_combinations):
            key = f"{abs_model}_{dist_model}_{inhib_model}"
            ax = axes[idx]

            if key in results_dict:
                result = results_dict[key]

                # Plot concentration-time curve
                ax.plot(result['time_points'], result['elimination']['plasma_concentration'],
                        color=colors[idx], linewidth=2)

                # Add PK parameters as text
                metrics = self.calculate_pk_metrics(
                    result['time_points'],
                    result['elimination']['plasma_concentration']
                )

                pk_text = f'C_max: {metrics["C_max"]:.3f} mg/L\n'
                pk_text += f'T_max: {metrics["T_max"]:.1f} h\n'
                pk_text += f'AUC: {metrics["AUC"]:.1f} mg·h/L'
                if metrics['t_half'] < 100:
                    pk_text += f'\nHalf-life: {metrics["t_half"]:.1f} h'

                ax.text(0.05, 0.95, pk_text, transform=ax.transAxes,
                        verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

                ax.set_xlabel('Time (h)')
                ax.set_ylabel('Plasma Concentration (mg/L)')
                ax.set_title(f'{abs_model.title()} Absorption + {dist_model.replace("_", " ").title()} Distribution')
                ax.grid(True, alpha=0.3)
            else:
                ax.text(0.5, 0.5, 'No Data', ha='center', va='center', fontsize=12)
                ax.set_xlabel('Time (h)')
                ax.set_ylabel('Plasma Concentration (mg/L)')
                ax.set_title(f'{abs_model.title()} Absorption + {dist_model.replace("_", " ").title()} Distribution')

        plt.tight_layout()
        # plt.savefig('../output/individual_inhibition_pbpk_comparison.png', dpi=300, bbox_inches='tight')
        print(f"Plot saved to: ../output/individual_inhibition_pbpk_comparison.png")
        plt.show()

    def plot_combined_comparison(self, results_dict):
        """Plot combined comparison"""
        plt.figure(figsize=(12, 8))

        absorption_models = ['first_order', 'adam']
        distribution_models = ['minimal_pbpk', 'full_pbpk']
        inhibition_models = [Inhibition_Model.COMPETITIVE_INHIBITION]

        # All models comparison
        plt.subplot(2, 2, 1)
        for key, result in results_dict.items():
            plt.plot(result['time_points'], result['elimination']['plasma_concentration'],
                     linewidth=2, label=key.replace('_', ' + '))
        plt.xlabel('Time (h)')
        plt.ylabel('Plasma Concentration (mg/L)')
        plt.title('All Models Comparison')
        plt.grid(True, alpha=0.3)
        plt.legend()

        # Absorption model comparison
        plt.subplot(2, 2, 2)
        for abs_model in absorption_models:
            for dist_model in distribution_models:
                for inhib_model in inhibition_models:
                    key = f"{abs_model}_{dist_model}_{inhib_model}"
                    if key in results_dict:
                        result = results_dict[key]
                        label = f"{abs_model} (all distribution)"
                        plt.plot(result['time_points'], result['elimination']['plasma_concentration'],
                                linewidth=2, label=label, alpha=0.7)
        plt.xlabel('Time (h)')
        plt.ylabel('Plasma Concentration (mg/L)')
        plt.title('Absorption Model Comparison')
        plt.grid(True, alpha=0.3)
        plt.legend()

        # Distribution model comparison
        plt.subplot(2, 2, 3)
        for abs_model in absorption_models:
            for dist_model in distribution_models:
                for inhib_model in inhibition_models:
                        key = f"{abs_model}_{dist_model}_{inhib_model}"
                        if key in results_dict:
                            result = results_dict[key]
                            label = f"{dist_model.replace('_', ' ')} (all absorption)"
                            plt.plot(result['time_points'], result['elimination']['plasma_concentration'],
                                    linewidth=2, label=label, alpha=0.7)
        plt.xlabel('Time (h)')
        plt.ylabel('Plasma Concentration (mg/L)')
        plt.title('Distribution Model Comparison')
        plt.grid(True, alpha=0.3)
        plt.legend()

        # PK parameters comparison table
        plt.subplot(2, 2, 4)
        ax = plt.gca()
        ax.axis('off')

        # Prepare table data
        table_data = []
        headers = ['Model', 'Cmax (mg/L)', 'Tmax (h)', 'AUC (mg·h/L)', 'Half-life (h)']

        for key, result in results_dict.items():
            metrics = self.calculate_pk_metrics(
                result['time_points'],
                result['elimination']['plasma_concentration']
            )
            table_data.append([
                key.replace('_', ' + '),
                f"{metrics['Cmax']:.3f}",
                f"{metrics['Tmax']:.1f}",
                f"{metrics['AUC']:.1f}",
                f"{metrics['Half_life']:.1f}" if metrics['Half_life'] < 100 else "N/A"
            ])

        if table_data:
            table = ax.table(cellText=table_data, colLabels=headers,
                             cellLoc='center', loc='center')
            table.auto_set_font_size(False)
            table.set_fontsize(9)
            table.scale(1.2, 1.5)
            plt.title('Pharmacokinetic Parameters Comparison')
        else:
            plt.text(0.5, 0.5, 'No valid data', ha='center', va='center', fontsize=12)

        plt.tight_layout()
        # plt.savefig('../output/individual_inhibition_pbpk_combined.png', dpi=300, bbox_inches='tight')
        print(f"Plot saved to: ../output/individual_inhibition_pbpk_combined.png")
        plt.show()


def main():
    """Main function"""
    print("=" * 60)
    print("Individual PBPK Modeling System")
    print("=" * 60)

    # 创建模拟器
    simulator = IndividualPBPKSimulator()

    # 定义要运行的模型组合
    absorption_models = ['first_order', 'adam']
    distribution_models = ['minimal_pbpk', 'full_pbpk']
    inhibition_models = [Inhibition_Model.COMPETITIVE_INHIBITION]

    # 存储所有结果
    all_results = {}

    # 运行所有模型组合
    for abs_model in absorption_models:
        for dist_model in distribution_models:
            for inhib_model in inhibition_models:
                key = f"{abs_model}_{dist_model}_{inhib_model}"

                try:
                    result = simulator.run_complete_pbpk(abs_model, dist_model, inhib_model)
                    all_results[key] = result

                    # 打印PK参数
                    metrics = simulator.calculate_pk_metrics(
                        result['time_points'],
                        result['elimination']['plasma_concentration']
                    )

                    print(f"\n  PK parameters:")
                    print(f"    C_max: {metrics['C_max']:.3f} mg/L")
                    print(f"    T_max: {metrics['T_max']:.1f} h")
                    print(f"    AUC: {metrics['AUC']:.1f} mg·h/L")
                    print(f"    Half-life: {metrics['t_half']:.1f} h" if metrics['t_half'] < 100 else "    Half-life: N/A")

                except Exception as e:
                    print(f"  Error: {str(e)}")
                    continue

    # Generate plots
    print("\nGenerating plots...")
    simulator.plot_all_models(all_results)
    # simulator.plot_combined_comparison(all_results)

    print("\nSimulation completed! Plots saved as:")
    print("  - individual_inhibition_pbpk_comparison.png")
    # print("  - individual_pbpk_combined.png")


if __name__ == "__main__":
    main()